{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Notebook Setup \n",
    "The following cell will install Drake, checkout the underactuated repository, and set up the path (only if necessary).\n",
    "- On Google's Colaboratory, this **will take approximately two minutes** on the first time it runs (to provision the machine), but should only need to reinstall once every 12 hours.  Colab will ask you to \"Reset all runtimes\"; say no to save yourself the reinstall.\n",
    "- On Binder, the machines should already be provisioned by the time you can run this; it should return (almost) instantly.\n",
    "\n",
    "More details are available [here](http://underactuated.mit.edu/drake.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "  import pydrake\n",
    "  import underactuated\n",
    "except ImportError:\n",
    "  !curl -s https://raw.githubusercontent.com/RussTedrake/underactuated/master/scripts/setup/jupyter_setup.py > jupyter_setup.py\n",
    "  from jupyter_setup import setup_underactuated\n",
    "  setup_underactuated()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Finding the limit cycle of the Van der Pol oscillator\n",
    "\n",
    "by setting up a simple trajectory optimization, where the timestep, $h$, is a decision variables."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "from underactuated.jupyter import SetupMatplotlibBackend\n",
    "plt_is_interactive = SetupMatplotlibBackend()\n",
    "\n",
    "from pydrake.examples.van_der_pol import VanDerPolOscillator\n",
    "from pydrake.all import DirectCollocation, PiecewisePolynomial, Solve\n",
    "\n",
    "plant = VanDerPolOscillator()\n",
    "context = plant.CreateDefaultContext()\n",
    "\n",
    "dircol = DirectCollocation(plant,\n",
    "                           context,\n",
    "                           num_time_samples=61,\n",
    "                           minimum_timestep=0.01,\n",
    "                           maximum_timestep=0.5)\n",
    "\n",
    "# Constrain all timesteps, $h[k]$, to be equal, so the trajectory breaks are evenly distributed.\n",
    "dircol.AddEqualTimeIntervalsConstraints()\n",
    "\n",
    "# Initial state on the surface of section (and velocity > .1).\n",
    "dircol.AddBoundingBoxConstraint([0., 0.1], [0., 10.], dircol.initial_state())\n",
    "\n",
    "# Periodicity constraint.\n",
    "# TODO(russt): Replace this with the vectorized version pending drake #8315.\n",
    "dircol.AddLinearConstraint(dircol.final_state()[0] == dircol.initial_state()[0])\n",
    "dircol.AddLinearConstraint(dircol.final_state()[1] == dircol.initial_state()[1])\n",
    "\n",
    "# Help the solver with an initial guess (circular trajectory).\n",
    "samples = np.linspace(0, 2 * np.pi, 10)\n",
    "x_guess = np.vstack(\n",
    "    ([2 * np.sin(t) for t in samples], [2 * np.cos(t) for t in samples]))\n",
    "initial_x_trajectory = PiecewisePolynomial.FirstOrderHold(samples, x_guess)\n",
    "\n",
    "dircol.SetInitialTrajectory(PiecewisePolynomial(), initial_x_trajectory)\n",
    "\n",
    "fig = plt.figure()\n",
    "h, = plt.plot([], [], \".-\")\n",
    "plt.xlim((-2.5, 2.5))\n",
    "plt.ylim((-3., 3.))\n",
    "\n",
    "def draw_trajectory(t, x):\n",
    "    h.set_xdata(x[0, :])\n",
    "    h.set_ydata(x[1, :])\n",
    "    fig.canvas.draw()\n",
    "    if plt.get_backend() == u\"MacOSX\":\n",
    "        plt.pause(1e-10)\n",
    "\n",
    "if plt_is_interactive:\n",
    "    dircol.AddStateTrajectoryCallback(draw_trajectory)\n",
    "\n",
    "result = Solve(dircol)\n",
    "assert result.is_success()\n",
    "\n",
    "x_trajectory = dircol.ReconstructStateTrajectory(result)\n",
    "\n",
    "x_knots = np.hstack([\n",
    "    x_trajectory.value(t) for t in np.linspace(x_trajectory.start_time(),\n",
    "                                               x_trajectory.end_time(), 100)\n",
    "])\n",
    "plt.plot(x_knots[0, :], x_knots[1, :]);\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
